/*
 * Copyright 2020 Makani Technologies LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#if !defined(__linux__)
#error "This AIO implementation should only be used for Linux targets."
#endif

#include "avionics/linux/aio.h"

#include <assert.h>
#include <errno.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <time.h>
#include <unistd.h>

#define __USE_MISC  // For struct ip_mreq.
#include <netinet/in.h>
#include <sys/socket.h>

#include "avionics/common/aio_header.h"
#include "avionics/common/aio_version_code.h"
#include "avionics/common/cvt.h"
#include "avionics/common/endian.h"
#include "avionics/common/network_config.h"
#include "avionics/common/pack_aio_header.h"
#include "avionics/linux/aio_socket.h"
#include "avionics/network/aio_node.h"
#include "avionics/network/message_type.h"
#include "common/macros.h"

// Configuration for an AIO connection.
// Main communication socket.
static int32_t g_fd;
// Connection identity.
static AioNode g_node;
// TODO: Move configuration to application, make constant.
static AioSocketConfig g_config;

// AIO message counters.
static int32_t g_received_valid_aio_packets = 0;
static int32_t g_received_invalid_aio_packets = 0;
static int32_t g_sent_aio_packets = 0;
// Most recent count of microseconds spent waiting for the user function
// or new messages.
static int64_t g_idle_usec = 0;

// Timer used to detect loss of link.
static int64_t g_last_packet_timestamp = 0;
// Maximum gap between received packets before we reset the socket.
static int64_t g_max_packet_gap_us = 0;
// Reset if we go for 5 Aio periods without a packet.
#define MAX_PERIODS_WITHOUT_PACKET 5
// Don't start looking for link loss until we've been up for 10 seconds.
#define MINIMUM_STARTUP_TIME_US (10 * 1000 * 1000)

// The current library has the following error handling behavior:
//   - Invalid arguments passed by the user are caught by assertions.
//   - Both a return code and an assertion are used if a there is a problem
//     setting up.
//   - Invalid data from the network is indicated by return codes.
// TODO: Add logging functionality to indicate the type of errors
// generated by AioSend or AioRecvToCvt.

int32_t AioSetup(AioNode node, uint16_t port,
                 const MessageType *subscribe_types,
                 int32_t num_subscribe_types) {
  assert(0 <= node && node < kNumAioNodes);
  assert(subscribe_types != NULL || num_subscribe_types == 0);
  assert(num_subscribe_types >= 0 && num_subscribe_types <= kNumMessageTypes);

  // Store source for outgoing messages.
  g_node = node;

  // Define configuration for AioSocket interface.
  AioSocketInitConfig(num_subscribe_types, subscribe_types, &g_config);
  AioSocketSetPort(port, &g_config);
  AioSocketSetLoopback(true, &g_config);

  // Allow for multiple listeners on the same port if we're not on the Q7.
  // TODO: Determine the rule for when to allow reusing the port.
#if !defined(__arm__)
  AioSocketSetReuseSocket(true, &g_config);  // Allow multiple processes.
#endif

  // Open AIO socket.
  g_fd = AioSocketOpen(&g_config);
  if (g_fd < 0) {
    assert(false);
    return -1;
  }
  return 0;
}

int32_t AioClose(void) {
  int32_t fd = g_fd;
  g_fd = -1;
  return AioSocketClose(fd);
}

static int32_t AioReopen(void) {
  int rc = AioClose();
  if (rc < 0) {
    return rc;
  }
  g_fd = AioSocketOpen(&g_config);
  if (g_fd < 0) {
    assert(false);
    return -1;
  }
  return 0;
}

// Waits up to timeout_us microseconds for a message to arrive.
// Returns -1 on an error, 1 if a message has arrived, and 0 if the
// timeout has elapsed.
static int32_t AioWaitForMessage(int64_t timeout_us) {
  // Block until next wakeup or message arrival.
  assert(g_fd >= 0);
  fd_set fds;
  FD_ZERO(&fds);
  FD_SET(g_fd, &fds);
  struct timeval tv = {(int32_t)(timeout_us / 1000000),
                       (int32_t)(timeout_us % 1000000)};

  // TODO: Ensure that select's timeout handling is monotonic.
  int32_t rc = select(g_fd + 1, &fds, NULL, NULL, &tv);
  assert(rc == 0 || rc == 1);
  if (rc == 0 &&
      (ClockGetUs() - g_last_packet_timestamp > g_max_packet_gap_us)) {
    if (AioReopen() == 0) {
      g_last_packet_timestamp = ClockGetUs();
    }
  }
  return rc;
}

// Attempts to handle a message from the socket.  If source or type
// are not NULL and a valid AIO message has been received without
// error, they will be populated with the received message's source
// and/or type.  If no message is received, or an error occurs -1 will
// be returned.  Otherwise, zero is returned.
static int32_t AioHandleMessage(AioNode *source, MessageType *type) {
  // Receive and sanity-check message.
  static uint8_t message[PACK_AIOHEADER_SIZE + AIO_MAX_MESSAGE_LEN];
  int32_t rc = AioSocketReceive(g_fd, sizeof(message), message);
  if (rc <= 0) {
    return -1;
  }

  // Get message receive time.
  // TODO: It would be more accurate to use the SIOCGSTAMP ioctl, but
  // the returned timestamp is not monotonic.
  int64_t timestamp = ClockGetUs();
  g_last_packet_timestamp = timestamp;

  // Validate there are enough bytes for the header to be processed.
  int32_t data_len = rc - PACK_AIOHEADER_SIZE;
  if (data_len < 0) {
    ++g_received_invalid_aio_packets;
    return -1;
  }
  AioHeader header;
  size_t unpacked_length = UnpackAioHeader(message, 1, &header);
  if (unpacked_length != PACK_AIOHEADER_SIZE || header.version != AIO_VERSION ||
      0 > header.source || header.source >= kNumAioNodes ||
      header.type >= kNumMessageTypes) {
    ++g_received_invalid_aio_packets;
    return -1;
  }
  ++g_received_valid_aio_packets;

  if (g_config.subscribed[header.type]) {
    if (type != NULL) {
      *type = header.type;
    }
    if (source != NULL) {
      *source = header.source;
    }
    CvtPut(header.source, header.type, header.sequence,
           timestamp, &message[PACK_AIOHEADER_SIZE], data_len);
    return 1;
  }

  return 0;
}

int32_t AioRecvToCvt(int64_t timeout_us, AioNode *source, MessageType *type) {
  int32_t rc = AioWaitForMessage(timeout_us);
  if (rc != 1) {
    return -1;
  } else {
    return AioHandleMessage(source, type);
  }
}

int64_t AioGetIdleUsec(void) {
  return g_idle_usec;
}

int32_t AioSendInternal(MessageType type, const void *buf, int32_t len,
                        const AioHeader *header) {
  if (buf == NULL || type >= kNumMessageTypes
      || len < 0 || len > (int32_t)AIO_MAX_MESSAGE_LEN) {
    assert(false);
    return -1;
  }

  // Prepare message.
  static uint8_t message[PACK_AIOHEADER_SIZE + AIO_MAX_MESSAGE_LEN];
  size_t packed_length = PackAioHeader(header, 1, message);
  if (packed_length != PACK_AIOHEADER_SIZE) {
    assert(false);
    return -1;
  }
  memcpy(&message[PACK_AIOHEADER_SIZE], buf, (size_t)len);

  // Send message.
  int32_t rc = AioSocketSend(&g_config, type, PACK_AIOHEADER_SIZE + len,
                             message, g_fd);
  g_sent_aio_packets += (rc > 0);
  return rc == -1 ? -1 : 0;
}

int32_t AioSend(MessageType type, const void *buf, int32_t len) {
  if (type >= kNumMessageTypes) {
    assert(false);
    return -1;
  }

  // Sequence numbers for outgoing messages.
  static uint16_t sequence_numbers[kNumMessageTypes];

  // Prepare header.
  AioHeader header = {
    .version = AIO_VERSION,
    .source = g_node,
    .type = type,
    .sequence = ++sequence_numbers[type],
    .timestamp = (uint32_t)ClockGetUs()
  };

  return AioSendInternal(type, buf, len, &header);
}

int32_t AioLoopStart(AioNode node, uint16_t port,
                     const MessageType *subscribe_types,
                     int32_t num_subscribe_types,
                     bool (*func)(void *), void *arg,
                     int64_t period_us) {
  assert(func != NULL);
  assert(period_us >= 0);

  int32_t rc = AioSetup(node, port, subscribe_types, num_subscribe_types);
  if (rc == -1) {
    assert(false);
    return -1;
  }

  // Compute initial wakeup time.
  int64_t now = ClockGetUs();
  int64_t wakeup = now + period_us;

  // Reset statistics.
  g_received_valid_aio_packets = 0;
  g_received_invalid_aio_packets = 0;
  g_sent_aio_packets = 0;

  // Prepare variables used for link-loss detection.
  g_last_packet_timestamp = now + MINIMUM_STARTUP_TIME_US;
  g_max_packet_gap_us = MAX_PERIODS_WITHOUT_PACKET * period_us;

  // Begin main loop.
  g_idle_usec = 0;
  while (true) {
    // Compute time to next wakeup.
    int64_t timeout = wakeup - ClockGetUs();

    // Run loop function.
    if (timeout < 0) {
      if (!func(arg)) return 0;
      wakeup += period_us;
      g_idle_usec = 0;
      continue;
    }

    int64_t wait_time = ClockGetUs();
    rc = AioWaitForMessage(timeout);
    g_idle_usec += ClockGetUs() - wait_time;

    if (rc == 1) AioHandleMessage(NULL, NULL);
  }

  return -1;
}

void AioGetStats(AioStats *stats) {
  stats->received_valid_aio_packets = (int16_t)g_received_valid_aio_packets;
  stats->received_invalid_aio_packets = (int16_t)g_received_invalid_aio_packets;
  stats->received_non_routine_packets = 0;
  stats->received_arp_request_packets = 0;
  stats->received_icmp_request_packets = 0;
  stats->received_probe_packets = 0;
  stats->sent_aio_packets = (int16_t)g_sent_aio_packets;

  g_received_valid_aio_packets = 0;
  g_received_invalid_aio_packets = 0;
  g_sent_aio_packets = 0;
}
